from __future__ import division, print_function
from svgpathtools import (
    Path,
    Line,
    CubicBezier,
    smoothed_path,
)
import svgwrite
import pandas as pd
import numpy as np
from svgwrite.params import Parameter
from svgwrite.base import BaseElement
from svgwrite.mixins import Transform
from svgwrite.container import SVG
from svgwrite.text import TSpan

# from consts import *
START_POS = (120, 0)
CAVAS_SIZE = (1650, 1250)
RATIO = 100.0
MARGIN = 20
FONT_SIZE = 12


class Mock(BaseElement, Transform):
    elementname = "g"
    _parameter = Parameter(True, "full")


class DrawEngine:
    def __init__(self, root_path="./", line_file="data.xlsx", output_dir="./") -> None:
        self.root_path = root_path
        self.line_file = line_file
        self.output_dir = output_dir

    def _reset(self, year) -> None:
        root_path = self.root_path
        line_file = self.line_file
        output_dir = self.output_dir

        self.hdf = pd.read_excel(f"{root_path}/data.xlsx", sheet_name="h_layer")
        self.vdf = pd.read_excel(f"{root_path}/data.xlsx", sheet_name="v_layer")
        node_df = pd.read_excel(f"{root_path}/data.xlsx", sheet_name="node")
        line_df = pd.read_excel(line_file, sheet_name="line")
        year = year if year in line_df.columns else int(year)
        line_df["val"] = line_df[year]
        line_df["svg_val"] = (line_df.val / RATIO).round(2)

        self.dwg = svgwrite.Drawing(
            f"{output_dir}/能流图_{year}.svg",
            size=CAVAS_SIZE,
            debug=True,
            style="background-color:white",
        )

        colors = line_df.color.values
        for i in range(len(colors)):
            if "," in str(colors[i]):
                colors[i] = (
                    f"rgb({colors[i]})" if "rgb(" not in str(colors[i]) else colors[i]
                )
        line_df["from_color"] = colors
        to_line_df = (
            line_df.groupby(by="from")
            .color.min()
            .reset_index()
            .rename(columns={"color": "to_color", "from": "to"})
        )

        line_df = line_df.merge(to_line_df, on="to", how="left")

        line_df.loc[line_df.to_color.isna(), "to_color"] = line_df.loc[
            line_df.to_color.isna(), "from_color"
        ]
        colors = []
        for _, row in line_df.iterrows():
            color = (
                row.from_color
                if row.from_color == row.to_color
                else self._gen_color([row.from_color, row.to_color])
            )
            colors.append(color)
        line_df["color"] = colors
        self.line_df = line_df

        colors = node_df.color.values
        for i in range(len(colors)):
            if "," in str(colors[i]):
                colors[i] = (
                    f"rgb({colors[i]})" if "rgb(" not in str(colors[i]) else colors[i]
                )

        node_df["color"] = colors
        left_node_df = (
            node_df[node_df.layer == 1]
            .merge(
                line_df[["from", "svg_val", "val"]],
                left_on="node",
                right_on="from",
            )
            .drop(columns=["from"])
        )

        right_node_df = node_df.merge(
            line_df.groupby(by="to")
            .agg({"svg_val": "sum", "val": "sum"})
            .reset_index(),
            left_on="node",
            right_on="to",
        ).drop(columns=["to"])

        node_df = pd.concat([left_node_df, right_node_df], ignore_index=True)
        node_df[["dy", "y1", "y2", "my", "y2e"]] = 0.0
        node_df["my"] = MARGIN
        end_layer = node_df.layer.max()
        node_df.loc[node_df.layer == end_layer, "my"] = MARGIN + 30.0
        node_df = node_df.merge(self.hdf, on="layer")
        self.node_df = node_df

        self.path_arr = []
        self.all_texts = []

    def update_input_nodes(self):
        node_df = self.node_df.copy()
        input_layer = node_df.layer.min()
        node_df["v"] = node_df.index.values
        index_df = (
            node_df[node_df.layer == input_layer].groupby("cat").v.min().reset_index()
        )
        index_df = index_df.merge(
            self.vdf[["layer", "dy"]], left_on="cat", right_on="layer"
        )

        self.node_df.loc[index_df.v.values, "dy"] = index_df.dy.values
        # self.node_df.loc[index_df.v.values, "my"] = 0
        data = []
        y2 = 0
        for index, row in self.node_df[self.node_df.layer == input_layer].iterrows():
            dy, my, v = row.dy, row.my, row.svg_val
            y1 = y2 + dy + my
            y2 = y1 + v
            data.append([index, y1, y2])
        data = np.array(data)
        self.node_df.loc[data[:, 0], ["y1", "y2"]] = data[:, 1:]

    def update_output_nodes(self):
        node_df = self.node_df.copy()
        end_layer = node_df.layer.max()
        node_df["v"] = node_df.index.values
        index_df = (
            node_df[node_df.layer == end_layer].groupby("cat").v.min().reset_index()
        )
        index_df = index_df.merge(
            self.vdf[["layer", "dy"]], left_on="cat", right_on="layer"
        )

        self.node_df.loc[index_df.v.values, "dy"] = index_df.dy.values

        data = []
        y2 = 0
        for index, row in self.node_df[self.node_df.layer == end_layer].iterrows():
            dy, my, v = row.dy, row.my, row.svg_val
            cnt = self.line_df[self.line_df["to"] == row.node].shape[0]
            y1 = y2 + dy + my
            y2 = y1 + v + 15 * (cnt - 1)
            data.append([index, y1, y2])
        data = np.array(data)
        print(f"y2:{y2}")
        self.max_y = y2
        self.node_df.loc[data[:, 0], ["y1", "y2"]] = data[:, 1:]

    def update_mid_nodes(self):
        for index, row in self.node_df.iterrows():
            if row.refer_node is None or row.refer_node == "":
                continue
            refer_nodes = self.node_df[self.node_df.node == row.refer_node]
            if refer_nodes.empty:
                continue
            refer_node = refer_nodes.iloc[0]

            y1 = refer_node.y1 if row.refer_direction == "top" else refer_node.y2
            y1 += row.refer_dy
            y2 = y1 + row.svg_val
            self.node_df.loc[index, ["y1", "y2"]] = [y1, y2]
        # print(self.node_df[self.node_df.node.isin(["洗练", "发电供热"])])

    def run_head_layer(self):
        input_layer = self.node_df.layer.min()
        node_df = self.node_df[self.node_df.layer == input_layer]
        node_df.to_clipboard(index=False)
        path_arr = []
        dwg = self.dwg
        for cat, frame in node_df.groupby("cat"):
            x = frame.iloc[0].x
            min_y, max_y = frame.y1.min(), frame.y2.max()
            y = (min_y + max_y + FONT_SIZE) / 2.0
            val = dwg.text(
                cat,
                insert=(x - 170, y),
                fill="black",
                style="font-weight:bolder",
            )
            dwg.add(val)

            p0 = complex(x - 90, min_y)
            p1 = complex(x - 100, p0.imag)
            p2 = complex(p1.real, max_y)
            p3 = complex(p0.real, p2.imag)
            seg0 = Line(p0, p1)
            seg1 = Line(p1, p2)
            seg2 = Line(p2, p3)
            left_path = Path(seg0, seg1, seg2)
            dwg.add(
                dwg.path(
                    d=left_path.d(),
                    stroke="black",
                    fill="none",
                    stroke_width=1,
                    stroke_linecap="round",
                )
            )
            for _, row in frame.iterrows():
                y = (row.y1 + row.y2) / 2.0
                val = dwg.text(
                    row.node,
                    insert=(x - 80, y),
                    fill="black",
                )
                dwg.add(val)

    def _gen_color(self, colors):
        dwg = self.dwg
        grad = dwg.linearGradient((0, 0), (1, 0))

        dwg.defs.add(grad)
        grad.add_stop_color(0, colors[0])
        grad.add_stop_color(1, colors[1])
        return grad.get_paint_server()

    def run_line_layer(self):
        dwg = self.dwg
        line_df = self.line_df
        node_df = self.node_df.copy()
        node_df["left_used"] = 0.0
        node_df["right_used"] = 0.0
        # node_df = node_df[node_df.layer <=2]

        all_paths = []
        all_colors = []
        all_vals = []
        all_texts = []
        end_layer = node_df.layer.max()
        input_layer = node_df.layer.min()
        for i, row in line_df.iterrows():
            if row.val == 0:
                continue

            from_nodes = node_df[node_df["node"] == row["from"]]
            end_nodes = node_df[node_df["node"] == row["to"]]

            if from_nodes.empty or end_nodes.empty:
                break
            color = row.color if "(" in row.color else f"hls({row.color})"
            all_colors.append(color)
            from_node, end_node = from_nodes.iloc[0], end_nodes.iloc[0].copy()

            x1, y1 = (
                from_node["x"],
                from_node["right_used"] + from_node["y1"] + row["svg_val"] / 2.0,
            )
            p1 = complex(x1, y1)

            my = 0.0 if end_node.layer != end_layer else 10.0

            x4, y4 = (
                end_node["x"] + 1,
                end_node["left_used"] + end_node["y1"] + +row["svg_val"] / 2.0,
            )

            if row.to == "氢能":
                y4 = y1
                node_df.loc[node_df.node == end_node.node, ["y1", "y2"]] = [
                    y4 - (end_node.y2 - end_node.y1) / 2.0,
                    y4 + (end_node.y2 - end_node.y1) / 2.0,
                ]

            if "#id" in row.color and y1 == y4:
                y4 += 0.001
            # if not is_calced:
            node_df.loc[node_df.node == end_node.node, "left_used"] += row.svg_val + my

            node_df.loc[node_df.node == from_node.node, "right_used"] += row.svg_val
            p4 = complex(x4, y4)

            if row["shape"] == "line":
                seg = Line(p1, p4)
                path = Path(seg)
                all_paths.append(path)
            else:
                p2 = complex(x4 + row.dx1 if row.refer1 == "to" else x1 + row.dx1, y1)
                p3 = complex(x4 + row.dx2, y4)
                seg0 = Line(p1, p2)
                seg1 = CubicBezier(
                    p2,
                    complex(p2.real / 2.0 + p3.real / 2.0, p2.imag),
                    complex(p2.real / 2.0 + p3.real / 2.0, p3.imag),
                    p3,
                )
                seg2 = Line(p3, p4)
                # print(seg0, seg1, seg2)
                path = Path(seg0, seg1, seg2)

                spath = smoothed_path(path, tightness=0.5, ignore_unfixable_kinks=True)
                all_paths.append(spath)
            all_vals.append(row.svg_val)

            text_x, text_y = x4 - 50, y4 + 6
            if from_node.layer == input_layer:
                text_x, text_y = x1 + 10, y1 + 6
            text = f"{np.round(row.val, 1)}"
            if row.to == "氢能":
                text = f"氢能 {text}"
            all_texts.append(
                dwg.text(
                    text,
                    insert=(text_x, text_y),
                    fill="black",
                    style="font-weight:bolder;font-size:10px",
                )
            )

        for path, color, val, text in zip(all_paths, all_colors, all_vals, all_texts):
            dwg.add(
                dwg.path(
                    d=path.d(),
                    stroke=color,
                    fill="none",
                    stroke_width=val,
                )
            )
            # if val > 2:
            #     dwg.add(text)
        self.all_texts += all_texts
        return node_df

    def run_end_layer(self):
        end_layer = self.node_df.layer.max()
        node_df = self.node_df[self.node_df.layer == end_layer]
        node_df.to_clipboard(index=False)
        dwg = self.dwg
        for _, row in node_df.iterrows():
            x = row.x
            min_y, max_y = row.y1, row.y2
            y = (min_y + max_y + FONT_SIZE) / 2.0

            p0 = complex(x + 10, min_y)
            p1 = complex(x + 15, p0.imag)
            p2 = complex(p1.real, max_y)
            p3 = complex(p0.real, p2.imag)
            seg0 = Line(p0, p1)
            seg1 = Line(p1, p2)
            seg2 = Line(p2, p3)
            left_path = Path(seg0, seg1, seg2)
            dwg.add(
                dwg.path(
                    d=left_path.d(),
                    stroke="black",
                    fill="none",
                    stroke_width=1,
                    stroke_linecap="round",
                )
            )

            text = dwg.text(
                row.node,
                insert=(x + 20, y),
                fill="black",
            )
            dwg.add(text)
        print(f"x:{x}")

    def run_loss_node(self, node_df):
        dwg = self.dwg
        node = node_df[node_df.node == "电力汇合点"].iloc[0]
        print(f"loss node:{node.svg_val - node.right_used}")
        if node.svg_val - node.right_used < 10:
            return
        line_df = self.line_df

        line_color = line_df[line_df.to == "电力汇合点"].iloc[0].from_color

        width = node.svg_val - node.right_used

        p0 = complex(node.x, node.y1 + node.right_used + width / 2.0 + 0.1)

        p1 = p0 + 5
        p2 = p1 + complex(width, 10)
        p3 = p2 + 10j
        p4 = complex(p3.real, self.max_y - 20)

        p_arr = [p0, p1, p2, p3, p4]

        color = self._gen_color([line_color, "rgb(138,137,132)"])
        seg_arr = [Line(p_arr[i], p_arr[i + 1]) for i in range(len(p_arr) - 1)]
        path = smoothed_path(Path(*seg_arr), tightness=0.88)
        dwg.add(
            dwg.path(
                d=path.d(),
                stroke=color,
                fill="none",
                stroke_width=width,
            )
        )
        dwg.add(
            dwg.text(
                str(np.round(node.val, 1)),
                insert=(p1.real + 20, p1.imag + 10),
                fill="black",
                style="font-weight:bolder;font-size:10px",
            )
        )

    def _draw_color_node(self, txt, val, color, cx, cy):
        dwg = self.dwg
        m = Mock()
        m.add(
            dwg.rect(
                size=(100, val),
                fill=color,
                stroke="black",
            )
        )
        sub = SVG(id="subsvg", size=(100, val))
        if len(txt) == 2:
            at = dwg.text(
                txt,
                insert=("50%", "50%"),
                font_size=12,
                dominant_baseline="middle",
                text_anchor="middle",
            )

        else:
            dx = [-7, 7] * 2
            dy = [-4, -4, 8, 8]
            x = ["50%"] * 4
            at = dwg.text(
                "",
                insert=("50%", "50%"),
                font_size=12,
                dominant_baseline="middle",
                text_anchor="middle",
            )
            at.add(TSpan(txt, x=x, y=x, dx=dx, dy=dy))
        sub.add(at)
        m.add(sub)

        m.translate(tx=cx, ty=cy)

        dwg.add(m)

    def draw_color_nodes(self):
        dwg = self.dwg
        node_df = self.node_df[~self.node_df.color.isnull()].reset_index(drop=True)
        for _, row in node_df.iterrows():
            node = row.node
            self._draw_color_node(node, row.y2 - row.y1, row.color, row.x, row.y1)

    def draw_assist_lines(self):
        dwg = self.dwg
        for _, layer in self.hdf.iterrows():
            dwg.add(
                dwg.line(
                    start=(layer.x * 1.0, 10),
                    end=(layer.x * 1.0, CAVAS_SIZE[1] - 30),
                    stroke="red",
                )
            )

    def draw_texts(self):
        dwg = self.dwg
        for text in self.all_texts:
            dwg.add(text)

    def run_by_year(self, year):
        self._reset(year)
        self.update_input_nodes()
        self.update_output_nodes()
        self.update_mid_nodes()

        self.run_head_layer()
        self.run_end_layer()

        node_df = self.run_line_layer()

        self.run_loss_node(node_df)

        self.draw_color_nodes()

        self.draw_texts()

        # self.draw_assist_lines()

        self.dwg.save(pretty=True)

    def get_years(self):
        line_df = pd.read_excel(self.line_file, sheet_name="line")
        columns = line_df.columns.values.astype(str)
        years = [col for col in columns if col[0] == "2" and len(col) == 4]
        return years

    def run(self):
        years = self.get_years()
        print(f"years:{years}")
        for year in years:
            self.run_by_year(year)
            # break


def get():
    import os

    return os.getcwd()


if __name__ == "__main__":
    engine = DrawEngine()
    engine.run()
